package com.example.wxg.stream.custom.partition;

import com.example.wxg.stream.custom.source.MyNoParallelSource;
import org.apache.flink.api.common.functions.MapFunction;
import org.apache.flink.api.common.functions.Partitioner;
import org.apache.flink.api.java.tuple.Tuple1;
import org.apache.flink.streaming.api.datastream.DataStream;
import org.apache.flink.streaming.api.datastream.DataStreamSource;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;

/**
 * @author void
 * @date 2021/11/12 14:10
 * @desc 自定义分区，奇偶分区
 */
public class MyPartition implements Partitioner<Long> {
    @Override
    public int partition(Long key, int i) {
        System.out.println("分区总数:"+i);
        if(key%2==0){
            return 0;
        }
        return 1;
    }

    public static void main(String[] args) throws Exception {
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        env.setParallelism(2);

        DataStreamSource<Long> text = env.addSource(new MyNoParallelSource());
        //数据转换 long->tuple1
        DataStream<Tuple1<Long>> tupleData = text.map(new MapFunction<Long, Tuple1<Long>>() {
            @Override
            public Tuple1<Long> map(Long aLong) throws Exception {
                return new Tuple1<>(aLong);
            }
        });
        
        //分区后的数据
        DataStream<Tuple1<Long>> partitionData = tupleData.partitionCustom(new MyPartition(), 0);
        DataStream<Long> result = partitionData.map(new MapFunction<Tuple1<Long>, Long>() {
            @Override
            public Long map(Tuple1<Long> value) throws Exception {
                System.out.println("当前线程id:"+Thread.currentThread().getId()+",value:"+ value);
                return value.getField(0);
            }
        });
        
        result.print().setParallelism(1);
        env.execute("启动自定义分区流");
    }
}
